import backtrader as bt
import datetime

class BTCNeutralGridStrategy(bt.Strategy):
    params = (
        ('grid_spacing', 100),  # 网格间距
        ('leverage', 5),  # 杠杆倍数
        ('initial_cash', 10000),  # 初始资金
        ('transaction_fee', 0.0004)  # 交易手续费率
    )

    def __init__(self):
        self.current_price = self.data.close
        self.position_size = 0
        self.balance = self.p.initial_cash
        self.prev_grid_level = None

    def next(self):
        current_price = self.current_price[0]
        # 计算当前网格区间
        grid_level = int(current_price / self.p.grid_spacing)

        # 如果是第一个交易日，不进行交易
        if self.prev_grid_level is None:
            self.prev_grid_level = grid_level
            return

        # 做多开仓逻辑
        if self.position.size == 0 and grid_level < self.prev_grid_level:
            # 计算开仓数量
            self.position_size = (self.balance * self.p.leverage) / current_price
            self.buy(size=self.position_size)
            # 扣除手续费
            self.balance -= self.position_size * current_price * (1 + self.p.transaction_fee)

        # 做空开仓逻辑
        elif self.position.size == 0 and grid_level > self.prev_grid_level:
            self.position_size = (self.balance * self.p.leverage) / current_price
            self.sell(size=self.position_size)
            # 扣除手续费
            self.balance -= self.position_size * current_price * (1 + self.p.transaction_fee)

        # 做多平仓逻辑
        elif self.position.size > 0 and grid_level >= self.prev_grid_level:
            self.sell(size=self.position_size)
            # 扣除手续费并更新余额
            self.balance += self.position_size * current_price * (1 - self.p.transaction_fee)
            self.position_size = 0

        # 做空平仓逻辑
        elif self.position.size < 0 and grid_level <= self.prev_grid_level:
            self.buy(size=-self.position_size)
            # 扣除手续费并更新余额
            self.balance += -self.position_size * current_price * (1 - self.p.transaction_fee)
            self.position_size = 0

        self.prev_grid_level = grid_level


if __name__ == '__main__':
    # 创建 Cerebro 引擎
    cerebro = bt.Cerebro()

    # 使用 YahooFinanceData 加载数据
    # 这里使用比特币 - USD（BTC - USD）的数据，你可以根据需要修改股票代码
    data = bt.feeds.YahooFinanceData(
        dataname='BTC-USD',
        fromdate=datetime.datetime(2023, 1, 1),
        todate=datetime.datetime(2023, 12, 31)
    )

    cerebro.adddata(data)

    # 添加策略
    cerebro.addstrategy(BTCNeutralGridStrategy)

    # 设置初始资金
    cerebro.broker.setcash(10000)

    # 运行回测
    print('初始资金: %.2f' % cerebro.broker.getvalue())
    cerebro.run()
    print('最终资金: %.2f' % cerebro.broker.getvalue())

    # 绘制结果
    cerebro.plot()